import apriltag
import cv2 as cv
import numpy as np
from uarm.wrapper import SwiftAPI
from math import sin, cos, pi

swift = SwiftAPI()
swift.reset()

at_detector = apriltag.Detector(apriltag.DetectorOptions(families='tag36h11'))

base2mark_pose = np.load("data/base2mark_pose.npz")["pose"]
print("base2mark_pose\n",base2mark_pose)

# 加载标定参数
with np.load('./data/camera_param.npz') as X:
    mtx, dist, _, _ = [X[i] for i in ('mtx','dist','rvecs','tvecs')]


def eulerAngles2RotationMatrix(theta):
    """
    欧拉角转换旋转矩阵
    :param theta: 欧拉角角度
    :return: 旋转矩阵
    """
    R_x = np.array([[1, 0, 0],
                    [0, cos(theta[0]), sin(theta[0])],
                    [0, -sin(theta[0]), cos(theta[0])]
                    ])
    R_y = np.array([[cos(theta[1]), 0, sin(theta[1])],
                    [0, 1, 0],
                    [-sin(theta[1]), 0, cos(theta[1])]
                    ])
    R_z = np.array([[cos(theta[2]), sin(theta[2]),0],
                    [-sin(theta[2]), cos(theta[2]),0],
                    [0, 0, 1]
                    ])
    R = R_z@R_y@R_x
    return R

def get_cam2base_T(pix_pos,obj_pos):
    """
    获取相机到世界坐标系的变换矩阵
    :param pix_pos: 像素坐标点，大于等于四个点
    :param obj_pos: 目标点世界坐标点列表，大于等于四个点
    :return: 相机到世界坐标系的变换矩阵
    """
    pix_pos = np.array(pix_pos, dtype=np.float32).reshape(-1, 1, 2)
    obj_pos = np.array(obj_pos, dtype=np.float32).reshape(-1, 3)
    ret, rvecs, tvecs = cv.solvePnP(obj_pos, pix_pos, mtx, dist)
    R = cv.Rodrigues(rvecs)[0]
    T = np.row_stack(((np.column_stack((R, tvecs))), [0, 0, 0, 1]))
    return T

def get_base2end_T():
    angle = swift.get_polar()[1]
    R = eulerAngles2RotationMatrix([0,0,angle-90])
    position = swift.get_position()
    T = np.row_stack(((np.column_stack((R, position))), [0, 0, 0, 1]))
    return T


def get_cam2end_T(num,image):
    gray = cv.cvtColor(image,cv.COLOR_BGR2GRAY)
    tags = at_detector.detect(gray)
    base2end_T = get_base2end_T()
    for tag in tags:
        if tag.tag_id == num:
            # print(tag.corners)

            cam2base_T = get_cam2base_T(tag.corners,base2mark_pose)
    cam2end_T = base2end_T@cam2base_T
    print(cam2end_T[:,3])
    return cam2end_T



if __name__ == "__main__":
    cap = cv.VideoCapture(4)

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        cv.imshow("frame",frame)
        key = cv.waitKey(10)
        if key == 27:
            break
        if key == ord("d"):
            swift.set_servo_detach()
        if key == ord("q"):
            cam2end_T = get_cam2end_T(2, frame)
            np.savez("./data/cam2end_T", T=cam2end_T)
    cv.destroyAllWindows()
    cap.release()













